﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch;
using TorchSharp.Modules;
using static Tensorboard.TensorShapeProto.Types;
using Tensorboard;
using NgpSharp.RayMaching;
using System.Diagnostics;

namespace NgpSharp
{
    public class NeRFRender : nn.Module<Tensor, Tensor, Tensor, Tuple<Tensor, Tensor>>
    {
        public float bg_radius { get; protected set; }
        public int num_step { get; set; } = 512;
        protected int density_scale;
        protected float min_near;
        protected float density_thresh;
        protected int mean_density = 0;
        protected int iter_density = 0;
        protected int mean_count = 0;
        protected int local_step = 0;

        public NeRFRender(string name, int density_scale = 1, 
            float min_near = 0.2f, float density_thresh = 0.01f, float bg_radius = -1) 
            : base(name)
        {
            this.density_scale = density_scale;
            this.min_near = min_near;
            this.density_thresh = density_thresh;
            this.bg_radius = bg_radius; //# radius of the background sphere.
        }

        public override Tuple<torch.Tensor, torch.Tensor> forward(torch.Tensor x, torch.Tensor d, Tensor aabb)
        {
            throw new NotImplementedException();
        }

        public virtual Dictionary<string, Tensor> Density(Tensor x, Tensor aabb)
        {
            throw new NotImplementedException();
        }

        public virtual Dictionary<string, Tensor> Color(Tensor x, Tensor d, Tensor? mask = null, Tensor? geo_feat = null)
        {
            throw new NotImplementedException();
        }

        public Dictionary<string, torch.Tensor> Render(Tensor rays_o, Tensor rays_d, Tensor aabb, bool staged = false, Tensor? bg_color = null, bool perturb = false, int max_ray_batch = 4096)
        {
            long B = rays_o.shape[0];
            long N = rays_o.shape[1];
            var device = rays_o.device;
            if(aabb.device != device) aabb = aabb.to(device);
            if (staged)
            {
                Tensor depth = torch.empty(B, N, device: device);
                Tensor image = torch.empty(B, N, 3, device: device);
                Dictionary<string, torch.Tensor> result = new();

                for (int b = 0; b < B; b++)
                {
                    long head = 0;
                    while (head < N)
                    {
                        using (var c = torch.NewDisposeScope())
                        {
                            long tail = Math.Min(head + max_ray_batch, N);
                            var r = RunOnesCuda(rays_o[TensorIndex.Slice(b, b + 1), TensorIndex.Slice(head, tail)],
                                rays_d[TensorIndex.Slice(b, b + 1), TensorIndex.Slice(head, tail)], aabb,
                                bg_color: bg_color, perturb: perturb, num_steps: num_step);
                            depth[TensorIndex.Slice(b, b + 1), TensorIndex.Slice(head, tail)] = r["depth"];
                            image[TensorIndex.Slice(b, b + 1), TensorIndex.Slice(head, tail)] = r["image"];
                            head += max_ray_batch;
                        }
                    }
                }
                result.Add("depth", depth);
                result.Add("image", image);
                return result;
            }
            else
            {
                var result = RunOnesCuda(rays_o, rays_d, aabb, num_steps: num_step, bg_color: bg_color, perturb: perturb);
                return result;
            }
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="rays_o"></param>
        /// <param name="rays_d"></param>
        /// <param name="num_steps"></param>
        /// <param name="upsample_steps"></param>
        /// <param name="bg_color"></param>
        /// <param name="perturb"></param>
        /// <returns></returns>
        private Dictionary<string, Tensor> RunOnes(Tensor rays_o, Tensor rays_d, Tensor aabb, int num_steps=128, 
            Tensor? bg_color = null, bool perturb = false)
        {
            //# rays_o, rays_d: [B, N, 3], assumes B == 1
            //# return: pred_rgb: [B, N, 3]

            using var c = torch.NewDisposeScope();

            long B = rays_o.shape[0];
            long N = rays_d.shape[1];
            var device = rays_o.device;

            long[] prefix = rays_o.shape[..^1];
            rays_o = rays_o.contiguous().view(-1, 3);
            rays_d = rays_d.contiguous().view(-1, 3);

            N = rays_o.shape[0]; //# N = B * N, in fact

            List<Tensor> nf = NearFarAABBFunction.apply(rays_o, rays_d, aabb, min_near);
            var nears = nf[0]; var fars = nf[1];
            nears.unsqueeze_(-1);
            fars.unsqueeze_(-1);

            Tensor z_vals = torch.linspace(0.0, 1.0, num_steps, device: device).unsqueeze(0); //# [1, T]
            z_vals = z_vals.expand([N, num_steps]); //# [N, T]
            z_vals = nears + (fars - nears) * z_vals; //# [N, T], in [nears, fars]

            //# perturb z_vals
            Tensor sample_dist = (fars - nears) / num_steps;
            if (perturb)
            {
                z_vals = z_vals + (torch.rand(z_vals.shape, device: device) - 0.5) * sample_dist;
            }

            //# generate xyzs
            Tensor xyzs = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * z_vals.unsqueeze(-1); //# [N, 1, 3] * [N, T, 1] -> [N, T, 3]
            xyzs = torch.min(torch.max(xyzs, aabb[TensorIndex.Slice(0, 3)]), aabb[TensorIndex.Slice(3)]); //# a manual clip.

            //# query SDF and RGB
            var density_outputs = Density(xyzs.reshape(-1, 3), aabb);
            foreach (var item in density_outputs)
                density_outputs[item.Key] = item.Value.view(N, num_steps, -1);

            //# upsample z_vals (nerf-like)
            Tensor deltas;
            Tensor alphas;
            Tensor alphas_shifted;
            Tensor weights;

            deltas = z_vals[TensorIndex.Ellipsis, TensorIndex.Slice(1)] - z_vals[TensorIndex.Ellipsis, TensorIndex.Slice(0, -1)]; //# [N, T+t-1]
            deltas = torch.cat([deltas, sample_dist * torch.ones_like(deltas[TensorIndex.Ellipsis, TensorIndex.Slice(0, 1)])], dim: -1);
            alphas = 1 - torch.exp(-deltas * density_scale * density_outputs["sigma"].squeeze(-1)); //# [N, T+t]
            alphas_shifted = torch.cat([torch.ones_like(alphas[TensorIndex.Ellipsis, TensorIndex.Slice(0, 1)]), 1 - alphas + 1e-15], dim: -1); //# [N, T+t+1]
            weights = alphas * torch.cumprod(alphas_shifted, dim: -1)[TensorIndex.Ellipsis, TensorIndex.Slice(0, -1)]; //# [N, T+t]

            Tensor dirs = rays_d.view(-1, 1, 3).expand_as(xyzs);
            foreach (var item in density_outputs)
                density_outputs[item.Key] = item.Value.view(-1, item.Value.shape[^1]);

            Tensor mask = weights > 1e-4; //# hard coded
            var cDic = Color(xyzs.reshape(-1, 3), dirs.reshape(-1, 3), mask: mask.reshape(-1), geo_feat: density_outputs["geo_feat"]);
            Tensor rgbs = cDic["color"].view(N, -1, 3); //# [N, T+t, 3]

            //# calculate weight_sum (mask)
            Tensor weights_sum = weights.sum(dim: -1); //# [N]

            //# calculate depth 
            Tensor ori_z_vals = ((z_vals - nears) / (fars - nears)).clamp(0, 1);
            Tensor depth = torch.sum(weights * ori_z_vals, dim: -1);

            //# calculate color
            Tensor image = torch.sum(weights.unsqueeze(-1) * rgbs, dim: -2); //# [N, 3], in [0, 1]

            //# mix background color
            if (bg_color is null)
                bg_color = 1;
            else bg_color = bg_color.reshape(-1, 3);

            image = image + (1 - weights_sum).unsqueeze(-1) * bg_color!;
            image = image.view([.. prefix, 3]);
            depth = depth.view(prefix);

            return new Dictionary<string, Tensor>()
            {
                { "depth", depth.MoveToOuterDisposeScope() },
                { "image", image.MoveToOuterDisposeScope() },
                { "weights_sum", weights_sum.MoveToOuterDisposeScope() }
            };
        }

        private Dictionary<string, Tensor> RunOnesCuda(Tensor rays_o, Tensor rays_d, Tensor aabb, int num_steps = 128,
            Tensor? bg_color = null, bool perturb = false)
        {
            using var c = torch.NewDisposeScope();

            long B = rays_o.shape[0];
            long N = rays_d.shape[1];
            var device = rays_o.device;

            long[] prefix = rays_o.shape[..^1];
            rays_o = rays_o.contiguous().view(-1, 3);
            rays_d = rays_d.contiguous().view(-1, 3);

            N = rays_o.shape[0];
            uint C = 3;
            Tensor nears = torch.empty([N, 1], dtype: rays_o.dtype, device: rays_o.device);
            Tensor fars = torch.empty([N, 1], dtype: rays_o.dtype, device: rays_o.device);
            Tensor xyzs = torch.empty([N, num_step, C], dtype: rays_o.dtype, device: rays_o.device);
            Tensor z_vals = torch.empty([N, num_step], dtype: rays_o.dtype, device: rays_o.device);
            Tensor deltas = torch.empty([N, num_step], dtype: rays_o.dtype, device: rays_o.device);
            RaysCoordsFunction.apply(rays_o, rays_d, aabb, nears, fars, xyzs, z_vals, deltas, min_near, perturb);

            var density_outputs = Density(xyzs.reshape(-1, 3), aabb);
            foreach (var item in density_outputs)
                density_outputs[item.Key] = item.Value.view(N, num_steps, -1);

            Tensor alphas = 1 - torch.exp(-deltas * density_scale * density_outputs["sigma"].squeeze(-1)); //# [N, T+t]
            Tensor alphas_shifted = torch.cat([torch.ones_like(alphas[TensorIndex.Ellipsis, TensorIndex.Slice(0, 1)]), 1 - alphas + 1e-15], dim: -1);
            Tensor weights = alphas * torch.cumprod(alphas_shifted, dim: -1)[TensorIndex.Ellipsis, TensorIndex.Slice(0, -1)]; //# [N, T+t]

            Tensor dirs = rays_d.view(-1, 1, 3).expand_as(xyzs);
            foreach (var item in density_outputs)
                density_outputs[item.Key] = item.Value.view(-1, item.Value.shape[^1]);
            Tensor mask = weights > 1e-4; //# hard coded
            var cDic = Color(xyzs.reshape(-1, 3), dirs.reshape(-1, 3), mask: mask.reshape(-1), geo_feat: density_outputs["geo_feat"]);
            Tensor rgbs = cDic["color"].view(N, -1, 3); //# [N, T+t, 3]

            //# calculate weight_sum (mask)
            Tensor weights_sum = weights.sum(dim: -1); //# [N]

            //# calculate depth 
            Tensor ori_z_vals = ((z_vals - nears) / (fars - nears)).clamp(0, 1);
            Tensor depth = torch.sum(weights * ori_z_vals, dim: -1);

            //# calculate color
            Tensor image = torch.sum(weights.unsqueeze(-1) * rgbs, dim: -2); //# [N, 3], in [0, 1]

            //# mix background color
            if (bg_color is null)
                bg_color = 1;
            else bg_color = bg_color.reshape(-1, 3);

            image = image + (1 - weights_sum).unsqueeze(-1) * bg_color!;
            image = image.view([.. prefix, 3]);
            depth = depth.view(prefix);

            return new Dictionary<string, Tensor>()
            {
                { "depth", depth.MoveToOuterDisposeScope() },
                { "image", image.MoveToOuterDisposeScope() },
                { "weights_sum", weights_sum.MoveToOuterDisposeScope() }
            };
        }
    }
}
